package com.xinziruo.richclient.engine.impl;

import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.util.Collection;
import java.util.Collections;

import org.apache.http.HeaderElement;
import org.apache.http.HeaderElementIterator;
import org.apache.http.HttpEntity;
import org.apache.http.HttpResponse;
import org.apache.http.client.HttpClient;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.client.methods.HttpUriRequest;
import org.apache.http.client.params.ClientPNames;
import org.apache.http.conn.ConnectionKeepAliveStrategy;
import org.apache.http.cookie.Cookie;
import org.apache.http.cookie.CookieOrigin;
import org.apache.http.cookie.CookieSpec;
import org.apache.http.cookie.CookieSpecFactory;
import org.apache.http.cookie.MalformedCookieException;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.DefaultHttpClient;
import org.apache.http.impl.cookie.BrowserCompatSpec;
import org.apache.http.message.BasicHeaderElementIterator;
import org.apache.http.params.BasicHttpParams;
import org.apache.http.params.HttpConnectionParams;
import org.apache.http.params.HttpParams;
import org.apache.http.protocol.HTTP;
import org.apache.http.protocol.HttpContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.xinziruo.richclient.engine.CrawlResponse;
import com.xinziruo.richclient.engine.IntrudeCrawlEngine;
import com.xinziruo.richclient.engine.Param;
import com.xinziruo.richclient.engine.PayloadSource;
import com.xinziruo.richclient.engine.impl.HttpClientCrawlResponse;
import com.xinziruo.richclient.util.ExceptionUtils;


/**
 * 
 * @author loudyn
 * 
 */
public class HttpClientCrawlEngine implements IntrudeCrawlEngine {

	private static Logger LOGGER = LoggerFactory.getLogger(HttpClientCrawlEngine.class);
	private final static String DEFAULT_CHARSET = "UTF-8";

	private String host;
	private String charset;
	private Param[] params;
	private PayloadSource payload;

	private int currentParam = 0;

	private final static PayloadSource NULL_PAYLOAD = new PayloadSource() {

		@Override
		public void reset() {
			// do nothing.
		}

		@Override
		public boolean hasNext() {
			// always return false.
			return false;
		}

		@Override
		public String getPayload() {
			return null;
		}
	};

	private HttpClient createHttpClient() {

		HttpParams params = new BasicHttpParams();
		HttpConnectionParams.setConnectionTimeout(params, 15000);
		HttpConnectionParams.setSoTimeout(params, 15000);

		CookieSpecFactory csf = new CookieSpecFactory() {
			public CookieSpec newInstance(HttpParams params) {
				return new BrowserCompatSpec() {

					@Override
					public void validate(Cookie cookie, CookieOrigin origin) throws MalformedCookieException {
					}
				};
			}
		};

		DefaultHttpClient client = new DefaultHttpClient(params);
		client.getCookieSpecs().register("ingore", csf);
		client.getParams().setParameter(ClientPNames.COOKIE_POLICY, "ingore");

		client.setKeepAliveStrategy(new ConnectionKeepAliveStrategy() {

			public long getKeepAliveDuration(HttpResponse response, HttpContext context) {
				// Honor 'keep-alive' header
				HeaderElementIterator it = new BasicHeaderElementIterator(response.headerIterator(HTTP.CONN_KEEP_ALIVE));
				while (it.hasNext()) {
					HeaderElement he = it.nextElement();
					String param = he.getName();
					String value = he.getValue();
					if (value != null && param.equalsIgnoreCase("timeout")) {
						try {

							return Long.parseLong(value) * 1000;
						} catch (NumberFormatException ignore) {
						}
					}
				}

				return 5 * 1000;

			}

		});

		return client;
	}

	public HttpClientCrawlEngine(String host) {
		this(host, DEFAULT_CHARSET, Collections.<Param> emptyList(), NULL_PAYLOAD);
	}

	public HttpClientCrawlEngine(HttpClient httpClient, String host) {
		this(host, DEFAULT_CHARSET, Collections.<Param> emptyList(), NULL_PAYLOAD);
	}

	public HttpClientCrawlEngine(String host, Collection<? extends Param> params, PayloadSource payload) {
		this(host, DEFAULT_CHARSET, params, payload);
	}

	public HttpClientCrawlEngine(String host, String charset, Collection<? extends Param> params, PayloadSource payload) {
		this.host = host;
		this.charset = charset;
		this.params = params.toArray(new Param[] {});
		this.payload = payload;
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see com.youboy.crawl.domain.intruder.IntruderEngine#hasNextRequest()
	 */
	@Override
	public boolean hasNextRequest() {
		if (currentParam >= params.length) {
			return false;
		}

		if (!params[currentParam].isAttack()) {
			currentParam++;
			return hasNextRequest();
		}

		if (!payload.hasNext()) {
			payload.reset();
			currentParam++;
			return hasNextRequest();
		}

		return true;
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see com.youboy.crawl.domain.intruder.IntruderEngine#issueRequest()
	 */
	@Override
	public CrawlResponse issueRequest() throws Exception {

		StringBuilder urlString = new StringBuilder();
		StringBuilder paramsString = new StringBuilder();
		StringBuilder bodyString = new StringBuilder();
		StringBuilder cookiesString = new StringBuilder();

		for (int i = 0; i < params.length; i++) {
			String value = (i == currentParam) ? payload.getPayload() : params[i].getValue();
			switch (params[i].getType()) {
			case URL:
				append(urlString, '/', params[i].getName().replaceAll("\\$\\{\\}", value), "", "", params[i].getEncode());
				break;
			case PARAM:
				append(paramsString, '&', params[i].getName(), "=", value, params[i].getEncode());
				break;
			case BODY:
				append(bodyString, '&', params[i].getName(), "=", value, params[i].getEncode());
				break;
			case COOKIE:
				append(cookiesString, ';', params[i].getName(), "=", value, params[i].getEncode());
				break;

			default:
				break;
			}
		}

		HttpUriRequest request = createHttpUriRequest(host, urlString, paramsString, bodyString, cookiesString);

		try {

			long startTime = System.currentTimeMillis();
			HttpResponse response = createHttpClient().execute(request);
			long elapsed = System.currentTimeMillis() - startTime;

			CrawlResponse crawlResp = new HttpClientCrawlResponse(response, elapsed);
			return crawlResp;
		} finally {
			request.abort();
		}
	}

	private HttpUriRequest createHttpUriRequest(String host,
												StringBuilder urlString,
												StringBuilder paramString,
												StringBuilder bodyString,
												StringBuilder cookiesString) throws UnsupportedEncodingException {

		String actualHost = createHost(host, urlString, paramString);
		actualHost = actualHost.replaceAll("\\n|\\r", "");
		if (bodyString.length() > 0) {
			HttpPost post = new HttpPost(actualHost);
			HttpEntity entity = new StringEntity(bodyString.toString(), charset);
			post.setEntity(entity);
			post.setHeader("Cookie", cookiesString.toString());
			return post;
		}

		HttpGet get = new HttpGet(actualHost);
		get.setHeader("Cookie", cookiesString.toString());

		LOGGER.warn("executing " + actualHost);
		return get;
	}

	private String createHost(String host, StringBuilder urlString, StringBuilder paramString) {

		if (urlString.length() <= 0 && paramString.length() <= 0) {
			return host;
		}

		StringBuilder builder = new StringBuilder().append(host);
		if(urlString.length() > 0){
			
			if (!host.endsWith("/")) {
				builder.append("/");
			}
			
			builder.append(urlString.toString());
		}
		
		if(paramString.length() > 0){
			
			if (host.indexOf('?') == -1) {
				builder.append('?');
			}
			
			builder.append(paramString.toString());
		}

		return builder.toString();
	}

	private void append(StringBuilder builder, char splitter, String name, String joinString, String value, String encode) {

		if (builder.length() > 0) {
			builder.append(splitter);
		}

		try {

			String encodeName = URLEncoder.encode(name, encode);
			String encodeValue = URLEncoder.encode(value, encode);
			builder.append(encodeName).append(joinString).append(encodeValue);
		} catch (UnsupportedEncodingException e) {
			throw ExceptionUtils.toUnchecked(e);
		}
	}
}
